#ifndef _2DOF_MODLE_H_
#define _2DOF_MODLE_H_

#include <pinocchio/parsers/urdf.hpp> // 引入 URDF 解析器
#include <pinocchio/algorithm/kinematics.hpp>
#include <pinocchio/algorithm/jacobian.hpp>
#include <pinocchio/algorithm/frames.hpp>
#include <pinocchio/algorithm/joint-configuration.hpp>
#include <pinocchio/multibody/model.hpp>
#include <pinocchio/algorithm/rnea.hpp>

#include <Eigen/Dense>
#include <math.h>
#include <stdio.h>
using namespace std;
using namespace Eigen;

class UR2DOFARM
{

public:
    // 改进DH
    float a[2] = {0}, d[2] = {0}, alpha[2] = {0}, offset[2] = {0};
    // 标准DH
    float ja[2] = {0}, jd[2] = {0}, jalpha[2] = {0}, joffset[2] = {0};

    // 改进 DH
    void SetDH(float set_a[2], float set_d[2], float set_alpha[2])
    {
        for (size_t i = 0; i < 2; i++)
        {
            a[i] = set_a[i];
            d[i] = set_d[i];
            alpha[i] = set_alpha[i];
        }
    }

    // 改进 DH
    void SetDH_SD(float set_a[2], float set_d[2], float set_alpha[2])
    {
        for (size_t i = 0; i < 2; i++)
        {
            ja[i] = set_a[i];
            jd[i] = set_d[i];
            jalpha[i] = set_alpha[i];
        }
    }

    // Vector2f Get_Tor(double q[2], double dq[2], double ddq[2])
    // {
    //     float m1, m2;
    //     Matrix3f R, R01, R12;
    //     Matrix3f I1, I2;
    //     Vector3f rci1, rci2, p1sta, p2sta, z;
    //     Vector3f omega00, omega11, omega22;
    //     Vector3f epsilon00, epsilon11, epsilon22;
    //     Vector3f a00, a11, a22;
    //     Vector3f ac11, ac22;
    //     Vector3f f00, f11, f22;
    //     Vector3f n00, n11, n22;
    //     Vector3f g0, g1, g2;
    //     Vector3f cros1, cros2;
    //     Vector2f tao;

    //     // 定义旋转矩阵
    //     R01 << cos(q[0]), -sin(q[0]), 0,
    //         sin(q[0]), cos(q[0]), 0,
    //         0, 0, 1;

    //     R12 << cos(q[1]), -sin(q[1]), 0,
    //         sin(q[1]), cos(q[1]), 0,
    //         0, 0, 1;

    //     // 定义惯量矩阵
    //     I1 << 0.00017, 0, 0,
    //         0, 0.00471, 0,
    //         0, 0, 0.00463;

    //     I2 << 0.00017, 0, 0,
    //         0, 0.04465, 0,
    //         0, 0, 0.04321;

    //     // 定义质量和质心
    //     m1 = 1.1;            // 连杆1的质量（包括电机）
    //     m2 = 1.132;          // 连杆2的质量（包括电机）
    //     rci1 << 0.12, 0, 0;  // 连杆1的质心位置（相对于连杆1坐标系）
    //     rci2 << 0.188, 0, 0; // 连杆2的质心位置（相对于连杆2坐标系）

    //     // 定义两个坐标系原点产生的向量
    //     p1sta << ja[0], 0, 0;
    //     p2sta << ja[1], 0, 0;

    //     // 初始化角速度、角加速度和线加速度
    //     omega00 << 0, 0, 0;
    //     epsilon00 << 0, 0, 0;
    //     a00 << 0, 0, 0;
    //     z << 0, 0, 1;

    //     // 计算角速度
    //     omega11 = R01.transpose() * (omega00 + z * dq[0]);
    //     omega22 = R12.transpose() * (omega11 + z * dq[1]);

    //     // 计算角加速度
    //     epsilon11 = R01.transpose() * (epsilon00 + omega00.cross(z * dq[0]) + z * ddq[0]);
    //     epsilon22 = R12.transpose() * (epsilon11 + omega11.cross(z * dq[1]) + z * ddq[1]);

    //     // 计算线加速度
    //     cros1 = omega11.cross(p1sta);
    //     a11 = R01.transpose() * a00 + epsilon11.cross(p1sta) + omega11.cross(cros1);
    //     cros1 = omega22.cross(p2sta);
    //     a22 = R12.transpose() * a11 + epsilon22.cross(p2sta) + omega22.cross(cros1);

    //     // 计算质心加速度
    //     cros1 = omega11.cross(rci1);
    //     ac11 = a11 + epsilon11.cross(rci1) + omega11.cross(cros1);
    //     cros1 = omega22.cross(rci2);
    //     ac22 = a22 + epsilon22.cross(rci2) + omega22.cross(cros1);

    //     // 计算重力
    //     g0 << 9.81, 0, 0;
    //     g1 = R01.transpose() * g0;
    //     g2 = (R01 * R12).transpose() * g0;

    //     // 初始化末端力和力矩
    //     f22 << 0.0, 0.0, 0.0;
    //     n22 << 0.0, 0.0, 0.0;

    //     // 计算力和力矩
    //     f11 = m1 * (ac11 - g1) + R12 * f22;
    //     f00 = m2 * (ac22 - g2) + f22;

    //     cros1 = I2 * omega22;
    //     cros2 = m2 * (ac22 - g2);
    //     n22 = I2 * epsilon22 + omega22.cross(cros1) + n22 + (p2sta.cross(f22)) + rci2.cross(cros2);

    //     cros1 = I1 * omega11;
    //     cros2 = m1 * (ac11 - g1);
    //     n11 = I1 * epsilon11 + omega11.cross(cros1) + R12 * n22 + (p1sta.cross(f11)) + rci1.cross(cros2);

    //     // 计算关节力矩
    //     tao(1) = (R12.transpose() * z).transpose() * n22;
    //     tao(0) = (R01.transpose() * z).transpose() * n11;

    //     return tao;
    // }

    Vector2f Get_Tor(const double q[2], const double dq[2], const double ddq[2],
                     pinocchio::Model &model, pinocchio::Data &data)
    {
        // 1. 设置重力方向（根据你的需求调整）
        // 默认是 [0, 0, -9.81]，你之前设置为 [9.81, 0, 0]（X轴方向）
        model.gravity.linear() << 9.81, 0, 0;

        // 2. 将输入数组转换为Eigen向量
        Vector2d q_vec(2); // 关节位置
        Vector2d v_vec(2); // 关节速度
        Vector2d a_vec(2); // 关节加速度

        q_vec << q[0], q[1];
        v_vec << dq[0], dq[1];
        a_vec << ddq[0], ddq[1];

        // 3. 检查模型维度是否匹配
        assert(model.nq == 2 && "Model dimension doesn't match input size!");

        // 4. 计算逆动力学（包含重力、科里奥利和惯性项）
        Vector2d tau = pinocchio::rnea(model, data, q_vec, v_vec, a_vec);
        return tau.cast<float>(); // 显式转换为float
    }

    // 阻抗控制函数
    void impedance_control_2dof(
        const Eigen::Vector2d &desired_end_pos,   // 期望的末端位置 (x, y)
        const Eigen::Vector2d &desired_end_vel,   // 期望的末端速度 (vx, vy)
        const Eigen::Vector2d &current_joint_pos, // 当前的关节位置 (q1, q2)
        const Eigen::Vector2d &current_joint_vel, // 当前的关节速度 (dq1, dq2)
        double stiffness,                         // 刚度系数 k
        double damping,                           // 阻尼系数 c
        pinocchio::Model &model,                  // Pinocchio 模型
        pinocchio::Data &data,                    // Pinocchio 数据
        Eigen::Vector2d &joint_torque             // 输出的关节力矩 (τ1, τ2)
    )
    {
        // 1. 计算当前末端位置和速度
        pinocchio::forwardKinematics(model, data, current_joint_pos);
        pinocchio::updateFramePlacements(model, data);

        pinocchio::FrameIndex ee_id = model.getFrameId("link2_end");                                                       // 末端执行器的 Frame ID
        Eigen::Vector3d ee_pos = data.oMf[ee_id].translation();                                                            // 当前末端位置 (x, y, z)
        Eigen::Vector3d ee_vel = pinocchio::getFrameVelocity(model, data, ee_id, pinocchio::LOCAL_WORLD_ALIGNED).linear(); // 当前末端速度 (vx, vy, vz)

        // 2. 计算末端误差
        Eigen::Vector2d position_error = desired_end_pos - ee_pos.head<2>(); // 位置误差 (ex, ey)
        Eigen::Vector2d velocity_error = desired_end_vel - ee_vel.head<2>(); // 速度误差 (dex, dey)

        // 3. 计算末端力 (阻抗控制律)
        Eigen::Vector2d end_force;
        end_force = stiffness * position_error + damping * velocity_error; // F = k * e + c * de

        // 4. 计算雅可比矩阵
        Eigen::MatrixXd J(6, model.nv); // 雅可比矩阵 (6x2)
        pinocchio::computeFrameJacobian(model, data, current_joint_pos, ee_id, pinocchio::LOCAL_WORLD_ALIGNED, J);

        // 5. 将末端力映射到关节力矩 (τ = J^T * F)
        Eigen::MatrixXd J_transposed = J.topRows<2>().transpose(); // 只取前两行（x, y 方向）
        joint_torque = J_transposed * end_force;                   // 关节力矩 (τ1, τ2)

        // 6. 重力补偿抵消重力
        double q[2] = {current_joint_pos(0), current_joint_pos(1)};
        double dq[2] = {current_joint_vel(0), current_joint_vel(1)};
        double ddq[2] = {0.0, 0.0}; // 假设加速度为 0
        Vector2f gravity_torque = Get_Tor(q, dq, ddq, model, data);

        joint_torque += gravity_torque.cast<double>(); // 从控制输出中减去重力
    }
};

#endif
